import numpy as np
import pandas as pd
from scipy.stats import uniform, weibull_min

def generate_population(
    male_distribution='Uniform',
    female_distribution='Uniform',
    male_ratio=0.5,
    n_samples=10,
    seed=42,
    low_age=0,
    up_age=100,
    alpha_male=1.5,
    beta_male=None,
    alpha_female=1.5,
    beta_female=None,
    p_below_65_male=0.5,  # 新增：男性Split分布时<65岁比例
    p_below_65_female=0.5 # 新增：女性Split分布时<65岁比例
):
    """
    生成人口数据，包括 ID、年龄和性别，允许为男性和女性指定不同的分布方式，
    并为 Split 分布分别指定低于 65 岁的比例。

    参数:
        male_distribution: 男性年龄分布方式 ('Uniform', 'Weibull', 'Split')
        female_distribution: 女性年龄分布方式 ('Uniform', 'Weibull', 'Split')
        male_ratio: 男性比例 (0 到 1)
        n_samples: 样本数量
        seed: 随机种子
        low_age: 年龄下限
        up_age: 年龄上限
        alpha_male: 男性 Weibull 分布的形状参数 (默认 1.5)
        beta_male: 男性 Weibull 分布的尺度参数 (默认 (up_age - low_age) / 1.5)
        alpha_female: 女性 Weibull 分布的形状参数 (默认 1.5)
        beta_female: 女性 Weibull 分布的尺度参数 (默认 (up_age - low_age) / 1.5)
        p_below_65_male: 男性在 Split 分布中年龄小于 65 岁的比例 (默认 0.5，仅在 male_distribution='Split' 时生效)
        p_below_65_female: 女性在 Split 分布中年龄小于 65 岁的比例 (默认 0.5，仅在 female_distribution='Split' 时生效)

    返回:
        DataFrame 包含 'id', 'age', 'gender' 列
    """
    np.random.seed(seed)
    if low_age >= up_age:
        raise ValueError("low_age must be less than up_age.")
    if not (0 <= male_ratio <= 1):
        raise ValueError("male_ratio must be between 0 and 1.")
    # 验证新的比例参数
    if not (0 <= p_below_65_male <= 1):
        raise ValueError("p_below_65_male must be between 0 and 1.")
    if not (0 <= p_below_65_female <= 1):
        raise ValueError("p_below_65_female must be between 0 and 1.")
    # 如果任何一个分布是 'Split'，则需要 65 岁在年龄范围内
    if (male_distribution == 'Split' or female_distribution == 'Split') and \
       (65 < low_age or 65 > up_age):
        raise ValueError("Split distribution requires 65 to be within the [low_age, up_age] range.")

    # 设置默认 beta
    if beta_male is None:
        # 根据Weibull分布在整个区间的合理性调整，可以考虑只在需要时计算
        beta_male = (up_age - low_age) / 1.5
    if beta_female is None:
        beta_female = (up_age - low_age) / 1.5

    # 计算男性和女性数量
    n_male = int(n_samples * male_ratio)
    n_female = n_samples - n_male

    # 定义年龄生成函数
    def generate_ages(n, distribution, gender):
        if n == 0: # 如果某性别数量为0，返回空数组
             return np.array([])

        if distribution == 'Uniform':
            return uniform.rvs(loc=low_age, scale=up_age - low_age, size=n)
        elif distribution == 'Weibull':
            alpha = alpha_male if gender == 'male' else alpha_female
            beta = beta_male if gender == 'male' else beta_female
            # 修正: Weibull 的尺度参数应该根据实际分布区间调整，这里简单处理，但可能需要更精细化
            # scale_adj = beta # (up_age - low_age) / 1.5 # 或者使用传入的beta
            return np.clip(
                weibull_min.rvs(c=alpha, loc=low_age, scale=beta, size=n),
                low_age,
                up_age
            )
        elif distribution == 'Split':
            # ***修改点：根据性别选择对应的 p_below_65***
            p_below = p_below_65_male if gender == 'male' else p_below_65_female
            n_below = int(n * p_below)
            n_above = n - n_below

            ages_below = np.array([])
            ages_above = np.array([])

            # 生成低于 65 岁的年龄（Uniform）
            if n_below > 0:
                 # 确保 low_age < 65
                 effective_up_age_below = min(65, up_age)
                 if low_age < effective_up_age_below:
                    ages_below = uniform.rvs(loc=low_age, scale=effective_up_age_below - low_age, size=n_below)
                 else: # 如果 low_age >= 65, 不能生成低于65的年龄
                     ages_below = np.full(n_below, low_age) # 或其他处理方式

            # 生成高于等于 65 岁的年龄（Weibull）
            if n_above > 0:
                alpha = alpha_male if gender == 'male' else alpha_female
                beta = beta_male if gender == 'male' else beta_female
                 # 确保 up_age > 65
                effective_low_age_above = max(65, low_age)
                if effective_low_age_above < up_age:
                    # 考虑Weibull的loc和scale在此区间的适应性
                    # scale_adj_above = beta # 可以根据(up_age - 65)调整beta，但这里用原始beta
                    ages_above = np.clip(
                        weibull_min.rvs(c=alpha, loc=effective_low_age_above, scale=beta, size=n_above),
                        effective_low_age_above,
                        up_age
                    )
                else: # 如果 up_age <= 65, 不能生成高于65的年龄
                    ages_above = np.full(n_above, up_age) # 或其他处理方式

            return np.concatenate([ages_below, ages_above])
        else:
            raise ValueError(f"Unsupported distribution '{distribution}'. Choose from 'Uniform', 'Weibull', 'Split'.")

    # 生成男性和女性的年龄
    male_ages = generate_ages(n_male, male_distribution, 'male')
    female_ages = generate_ages(n_female, female_distribution, 'female')

    # 如果某个年龄数组因为人数为0而是空，处理concatenate
    ages_list = []
    if len(male_ages) > 0:
        ages_list.append(male_ages)
    if len(female_ages) > 0:
        ages_list.append(female_ages)

    if not ages_list: # 如果总样本为0
        ages = np.array([])
        genders = np.array([])
        ids = np.array([])
    else:
        ages = np.concatenate(ages_list)
        # 生成性别
        genders = np.array(['male'] * n_male + ['female'] * n_female)
        # 生成 ID
        ids = np.arange(1, n_samples + 1)

    # 构建 DataFrame
    df = pd.DataFrame({
        'id': ids,
        'age': ages,
        'gender': genders
    })

    # 打乱顺序后再按 ID 排序，以混合男性和女性数据，而不是简单拼接
    df = df.sample(frac=1, random_state=seed).sort_values(by='id').reset_index(drop=True)

    return df
